package com.pinfly.common.crypto;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.DataOutputStream;
import java.io.IOException;
import java.io.ObjectInputStream;
import java.io.ObjectOutputStream;
import java.security.Key;
import java.security.MessageDigest;

import javax.crypto.Cipher;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;

import org.apache.log4j.Logger;

import com.pinfly.common.util.HexConvert;

/**
 * This class provides support for encrypting and decrypting data using AES-128.
 * The algorithm used ensures that encrypting the same plain text at a later
 * time, will produce different cipher text.
 * <p>
 * Note: In cryptography, the Advanced Encryption Standard (AES) is an
 * encryption standard adopted by the U.S. government. The standard comprises
 * three block ciphers, AES-128, AES-192 and AES-256. This class use AES-128.
 * </p>
 * <p>
 * Callers should use the CryptoFactory getInstance methods to obtain an
 * implementation.
 * </p>
 * 
 * 
 * This code was copied from a UMA project by Marc Lottman.
 * 
 */
public class AESEncryption implements ICrypto
{
    private static final Logger logger = Logger.getLogger (AESEncryption.class);

    private static final String ALGORITHM = "AES";
    private static final String CIPHER_TYPE = "AES/CBC/PKCS5Padding";
    private static final String DIGEST_TYPE = "SHA-256";

    private static final byte[] DEFAULT_SALT =
    { 60, 0, 35, 71, 1, 60, 32, 91, 122, 121, 67, 84, 79, 93, 52, 75 };
    private static final byte[] DEFAULT_IV_SUFFIX =
    { -7, -48, -54, -118, -5, 98, -11, -45 };
    private static final String UTF8 = "UTF-8";

    private static final byte[] s_keyData =
    { (byte) 0xF9, (byte) 98, (byte) 104, (byte) 79, (byte) 93, (byte) 63, (byte) 0xE1, (byte) 1, (byte) 117,
     (byte) 10, (byte) 37, (byte) 0xC4, (byte) 25, (byte) 7, (byte) 0xF5, (byte) 0xEE };

    protected Key m_key;

    /**
     * Creates an instance of the AESEncryption class using the default key text
     * and default salt.
     * 
     * Callers should use the CryptoFactory getInstance methods to obtain an
     * instance of this implementation.
     * 
     * @throws CryptoException if there are any problems obtaining the required
     *             encryption algorithms.
     */
    public AESEncryption () throws CryptoException
    {
        try
        {
            m_key = new SecretKeySpec (s_keyData, ALGORITHM);

            // Check that we can get the appropriate Cipher and digest.
            // We don't actually use or save them here, but want to error
            // quickly if for some strange reason the ciphers are not available.
            Cipher.getInstance (CIPHER_TYPE);
            MessageDigest.getInstance (DIGEST_TYPE);

        }
        catch (Exception e)
        {
            logger.warn (e.getMessage ());
            throw new CryptoException (e);
        }
    }

    public Key getKey ()
    {
        return m_key;
    }

    /**
     * Creates an instance of the AESEncryption class using the specified key
     * text and the default salt.
     * 
     * Callers should use the CryptoFactory getInstance methods to obtain an
     * instance of this implementation.
     * 
     * @param keyText The key used to encrypt the data.
     * @throws CryptoException if there are any problems obtaining the required
     *             encryption algorithms.
     * @throws IllegalArgumentException if keyText is null
     */
    public AESEncryption (String keyText) throws CryptoException
    {
        this (keyText, DEFAULT_SALT);
    }

    /**
     * Creates an instance of the AESEncryption class using the specified key
     * and salt.
     * 
     * Callers should use the CryptoFactory getInstance methods to obtain an
     * instance of this implementation.
     * 
     * @param keyText The key used to encrypt the text.
     * @param keySalt The salt used by the encryption/decryption.
     * @throws CryptographyException if there are any problems obtaining the
     *             required encryption algorithms.
     * @throws IllegalArgumentException if keyText or keySalt is null
     */
    public AESEncryption (String keyText, byte[] keySalt) throws CryptoException
    {

        if (keyText == null)
        {
            throw new IllegalArgumentException ("keyText cannot be null");
        }

        if (keySalt == null)
        {
            throw new IllegalArgumentException ("keySalt cannot be null");
        }

        try
        {
            // Derive a key from the key text and salt.
            // The idea here is to make it more computationally expensive to
            // test various keyTexts, since for each guess you would have to run
            // it
            // thru the below algorithm. In addition it allows us to create
            // the size of the key we need. Java provides a PBKDF2WithHmacSHA1
            // function, to derive the key, but unfortunately it is not
            // available in
            // IBM Java 5.
            MessageDigest md = MessageDigest.getInstance (DIGEST_TYPE);
            byte[] keyData = keyText.getBytes (UTF8);

            // Recursively hash - adding the hash result plus the salt
            // together each time.
            for (int i = 0; i < 16000; i++)
            {
                // normally one would create the data byte array outside
                // of the loop, but the size of it varies in the first
                // call since the initial keyData is just the bytes
                // from the keyText. The rest of the calls have a
                // keyData that is the size of the digest, which is static.
                byte[] data = new byte[keyData.length + keySalt.length];
                System.arraycopy (keyData, 0, data, 0, keyData.length);
                System.arraycopy (keySalt, 0, data, keyData.length, keySalt.length);
                keyData = md.digest (data);
            }

            // Create a secret key with the first 16 bytes (128 bits)
            // of the keyData.
            m_key = new SecretKeySpec (keyData, 0, 16, ALGORITHM);

            // Check that we can get the appropriate Cipher.
            // We don't actually use the Cipher here, but want to error
            // quickly if for some strange reason the cipher is not available.
            Cipher.getInstance (CIPHER_TYPE);

        }
        catch (Exception e)
        {
            logger.warn (e.getMessage ());
            throw new CryptoException (e);
        }
    }

    /**
     * Generate initialization vector for AES encryption.
     * 
     * @return Initialization Vector.
     */
    private static byte[] createIV ()
    {
        ByteArrayOutputStream s;
        try
        {
            s = new ByteArrayOutputStream ();
            DataOutputStream out = new DataOutputStream (s);
            out.writeLong (System.currentTimeMillis ());
            out.write (DEFAULT_IV_SUFFIX);
            out.flush ();
            out.close ();
            return s.toByteArray ();
        }
        catch (IOException e)
        {
            logger.warn (e.getMessage ());
            // exceptions are not raised on byte array output
            // streams, so this should never happen.
            // Just in case we convert it to a runtime exception
            throw new RuntimeException (e);
        }
    }

    public String decrypt (String encryptedData) throws CryptoException
    {
        if (encryptedData == null)
        {
            throw new IllegalArgumentException ("encryptedData cannot be null");
        }

        try
        {
            byte[] data = HexConvert.convertToBytes (encryptedData);
            byte[] decryptedData = decrypt (data);

            return new String (decryptedData);
        }
        catch (CryptoException e)
        {
            logger.warn (e.getMessage ());
            throw e;
        }
        catch (Exception e)
        {
            logger.warn (e.getMessage ());
            throw new CryptoException (e);
        }
    }

    public byte[] decrypt (byte[] encryptedData) throws CryptoException
    {
        if (encryptedData == null)
        {
            throw new IllegalArgumentException ("encryptedData cannot be null");
        }

        try
        {
            // pull off the initialization vector, it is on the front of the
            // data.

            if (encryptedData.length < 16)
            {
                throw new CryptoException ("Encrypted data does not contain necessary initialization vector");
            }
            byte[] iv = new byte[16];
            System.arraycopy (encryptedData, 0, iv, 0, 16);

            Cipher cipher = Cipher.getInstance (CIPHER_TYPE);
            IvParameterSpec ivSpec = new IvParameterSpec (iv);
            cipher.init (Cipher.DECRYPT_MODE, getKey (), ivSpec);

            byte[] dataWithHash = cipher.doFinal (encryptedData, 16, encryptedData.length - 16);
            if (dataWithHash.length < 32)
            {
                throw new CryptoException ("Encrypted data does not contain necessary hash");
            }

            byte[] savedDigest = new byte[32];
            System.arraycopy (dataWithHash, 0, savedDigest, 0, 32);

            byte[] decryptedData = new byte[dataWithHash.length - 32];
            System.arraycopy (dataWithHash, 32, decryptedData, 0, decryptedData.length);

            // Create our own digest to compare
            MessageDigest md = MessageDigest.getInstance (DIGEST_TYPE);
            byte[] digest = md.digest (decryptedData);

            if (!MessageDigest.isEqual (savedDigest, digest))
            {
                throw new CryptoException ("Encrypted data has been tampered with.  Hashes do not match.");
            }

            return decryptedData;
        }
        catch (CryptoException e)
        {
            logger.warn (e.getMessage ());
            throw e;
        }
        catch (Exception e)
        {
            logger.warn (e.getMessage ());
            throw new CryptoException (e);
        }
    }

    public Object decryptObject (byte[] data) throws CryptoException
    {
        if (data == null)
        {
            throw new IllegalArgumentException ("data cannot be null");
        }

        ObjectInputStream in = null;
        try
        {
            byte[] decrypted = decrypt (data);
            in = new ObjectInputStream (new ByteArrayInputStream (decrypted));

            return in.readObject ();
        }
        catch (Exception e)
        {
            logger.warn (e.getMessage ());
            throw new CryptoException (e);
        }
        finally
        {
            if (in != null)
            {
                try
                {
                    in.close ();
                }
                catch (IOException ignored)
                {
                }
            }
        }
    }

    public String encrypt (String data) throws CryptoException
    {
        if (data == null)
        {
            throw new IllegalArgumentException ("data cannot be null");
        }

        try
        {
            // We create a hash of the data 32 bytes, and pre-pend it to the
            // data to encrypt. This allows us to verify in the
            // decrypt that the data has not be changed in any way.
            byte[] bData = data.getBytes (UTF8);
            byte[] results = encrypt (bData);

            return HexConvert.convertHexString (results);
        }
        catch (Exception e)
        {
            logger.warn (e.getMessage ());
            throw new CryptoException (e);
        }
    }

    public byte[] encrypt (byte[] data) throws CryptoException
    {
        if (data == null)
        {
            throw new IllegalArgumentException ("data cannot be null");
        }

        try
        {
            // We create a hash of the data 32 bytes, and pre-pend it to the
            // data to encrypt. This allows us to verify in the
            // decrypt that the data has not be changed in any way.
            MessageDigest md = MessageDigest.getInstance (DIGEST_TYPE);
            byte[] digest = md.digest (data);
            byte[] iv = createIV ();
            byte[] dataWithHash = new byte[digest.length + data.length];
            System.arraycopy (digest, 0, dataWithHash, 0, digest.length);
            System.arraycopy (data, 0, dataWithHash, digest.length, data.length);

            Cipher cipher = Cipher.getInstance (CIPHER_TYPE);
            cipher.init (Cipher.ENCRYPT_MODE, getKey (), new IvParameterSpec (iv));

            byte[] encrypted = cipher.doFinal (dataWithHash);
            byte[] results = new byte[encrypted.length + 16];
            System.arraycopy (iv, 0, results, 0, 16);
            System.arraycopy (encrypted, 0, results, 16, encrypted.length);

            return results;
        }
        catch (Exception e)
        {
            logger.warn (e.getMessage ());
            throw new CryptoException (e);
        }
    }

    public byte[] encryptObject (Object obj) throws CryptoException
    {
        if (obj == null)
        {
            throw new IllegalArgumentException ("obj cannot be null");
        }

        ObjectOutputStream out = null;
        try
        {
            ByteArrayOutputStream bytes = new ByteArrayOutputStream ();
            out = new ObjectOutputStream (bytes);
            out.writeObject (obj);
            out.close ();

            byte[] retb = encrypt (bytes.toByteArray ());

            return retb;
        }
        catch (Exception e)
        {
            logger.warn (e.getMessage ());
            throw new CryptoException (e);
        }
        finally
        {
            if (out != null)
            {
                try
                {
                    out.close ();
                }
                catch (IOException ignored)
                {
                }
            }
        }
    }
}
